from src.utils import *


def tsne_2D(X: np.ndarray, Y: np.ndarray, title):
    X2 = TSNE(random_state=0).fit_transform(X)
    num_class = len(np.unique(Y))
    colors = sns.color_palette("hls", n_colors=num_class)
    # plt.figure(dpi=300)
    sns.scatterplot(
        x=X2[:, 0],
        y=X2[:, 1],
        hue=Y,
        palette=colors,
    )
    plt.title(title)


def tsne_2D_many_labels(X: np.ndarray, Y: pd.DataFrame):
    X2 = TSNE(random_state=0).fit_transform(X)
    n = len(Y.columns)
    scale = 5
    plt.figure(figsize=(n * 5, 5), dpi=100)

    for i, col in enumerate(Y.columns):
        ax = plt.subplot(1, n, i + 1)
        y = Y[col]
        num_class = len(np.unique(y))
        colors = sns.color_palette("hls", n_colors=num_class)
        # plt.figure(dpi=300)
        sns.scatterplot(
            x=X2[:, 0],
            y=X2[:, 1],
            hue=y,
            palette=colors,
            ax=ax,
        )
        plt.title(col)


def vis_tsne_train_Class():
    """
    train里面的所有变量作为X，train里面的Class（0，1）作为Y，可视化tsne
    """
    df = basic_fillna(train_csv).drop(columns="Id")
    print(df.isna().any().any())
    X, Y = df.drop(columns="Class"), df["Class"]
    X = pd.get_dummies(X)
    tsne_2D(X, Y, "train-Class")
    plt.show()


def vis_tsne_train_greeks(binarize_labels=False, EJ=True, Class=True):
    """
    train里面的所有变量作为X，greeks里面的Alpha,Beta,Gamma,Delta作为Y，可视化tsne
    """
    df = basic_fillna(train_csv).drop(columns="Id")
    print(df.isna().any().any())
    if not EJ:
        X = df.drop(columns="Class")
        greeks = greeks_csv.drop(columns=["Epsilon", "Id"])
        X = pd.get_dummies(X)
    else:
        X = df.drop(columns=["Class", "EJ"])
        greeks = greeks_csv.drop(columns=["Epsilon", "Id"])
        greeks['EJ'] = df['EJ']
    if Class:
        greeks['Class'] = df['Class']

    if binarize_labels:
        greeks = to_binary_labels(greeks)

    tsne_2D_many_labels(
        X=X,
        Y=greeks,
    )
    plt.show()
